#######################################################
# The class of general forests
# Author: moonho.hwang@gmail.com
#   Date: Oct.19.2011
#######################################################
from Graph import * # base class

# a constant string for edge type 

S_hier = 'hier' # keyword for hierarchical edge type

class Forest(Graph):
   """A class for a forest that is a set of trees."""
   # Constructor  
   def __init__(self):
      Graph.__init__(self)

   def parent(self,v):
      '''Get the parent of v in terms of hierarchy '''
      parents = self.inV(v, lambda G, edge: edge[S_type]==S_hier);
      if len(parents) == 0: return None
      elif len(parents) > 1: assert(0)
      else: return list(parents)[0]

   def setHierarchy(self, from_v, to_v):
      '''Set hierarchy from_v -['hier']->to_v '''
      self.addEdge(from_v, to_v, {S_type:S_hier})

   def addChild(self, parent, child):
      '''add a child <i>child</i> to the parent <i> parent </i>'''
      self.setHierarchy(parent, child)
      	
   def children(self, v):
      '''Get all children of v in terms of hierarchy. '''
      if v in self.V: 
        return self.outV(v, lambda G, e: e[S_type]==S_hier)
      else:
        return set() 

   def isRoot(self, v):
      '''Return true if v is a root in terms of hierarchy; false otherwise'''
      if v not in self.V: return False
      if self.parent(v) == None: return True
      else: return False

   def isLeaf(self, v):
      '''Return true if v is a leaf in terms of hierarchy; false otherwise '''
      if v not in self.V: return False # v is not registered in V yet
      if len(self.children(v)) == 0: return True
      else: return False
      
   def ancestors(self, v):
      '''Get the ancestors of v in terms of hierarchy '''
      return self.tinV(v, lambda G, e: e[S_type]==S_hier)
   
   def descendents(self, v):
      '''Get ancestors of v in terms of hierarchy '''
      return self.toutV(v, lambda G, e: e[S_type]==S_hier)
 
   def siblings(self, v):
      '''Get siblings which are children(parent(v)) - {v} '''
      return self.children(self.parent(v))-set([v]) # v itself is not incluced in the answer
#   end of Forest class

########### built-in testing codes ################
if __name__ == '__main__':
    f = Forest()
    print f.__doc__
    f.addVtx('1'); f.addVtx('1.1'); f.addVtx('1.2');
    f.addVtx('1.1.1'); f.addVtx('1.1.2');
    f.addVtx('1.2.1'); f.addVtx('1.2.2');
    
    f.setHierarchy('1','1.1'); f.setHierarchy('1','1.2');
    f.setHierarchy('1.1','1.1.1'); f.setHierarchy('1.1','1.1.2');
    f.setHierarchy('1.2','1.2.1'); f.setHierarchy('1.2','1.2.2');

    for v in f.V:
       print "%s:  parent=%s,  ancestors=%s,  Children=%s,  descendents=%s, siblings= %s."% \
       (str(v), str(f.parent(v)), str(f.ancestors(v)), str(f.children(v)), \
        str(f.descendents(v)), str(f.siblings(v)) )
